module Transform
  ( optimize
  , normalize
  ) where

import Prelude

import AST (Expr(..), Instr(..), Oper(..), Proc)
import Data.Bifunctor (bimap)
import Data.Foldable (foldMap)
import Data.FoldableWithIndex (foldMapWithIndex)
import Data.List (List(..), (:))
import Data.Map as M
import Data.Maybe (Maybe(..), maybe)
import Data.Ord.Max (Max(..))
import Data.Ord.Min (Min(..))
import Data.TraversableWithIndex (traverseWithIndex)
import Data.Tuple.Nested (type (/\))
import RList (RList(..), revToList, (:>))

type Pass = Proc -> Proc

type BigPass = RList Instr -> Proc -> RList Instr

run :: BigPass -> Pass
run pass = pass RNil >>> revToList

liftPass :: Pass -> BigPass
liftPass pass r proc = case pass proc of
  Loop x : xs -> liftPass pass (r :> Loop x') xs
    where
    x' = x { proc = run (liftPass pass) x.proc }
  x : xs -> liftPass pass (r :> x) xs
  Nil -> r

merge :: Pass
merge (Update m : Update n : xs) = merge (Update m' : xs)
  where
  m' = M.unionWith union m (inst <$> n)
  union (Add a) (Add b) = Add (a + b)
  union (Set a) (Add b) = Set $ EAdd a (ECst b)
  union _ a = a
  inst (Set a) = Set $ instE getExpr a
  inst a = a
  getExpr i = maybe (EDyn i) (operExpr i) (M.lookup i m)
merge xs = xs

simpl :: Pass
simpl (Move 0 : xs) = xs
simpl (Update m : xs) = xs'
  where
  xs' = if M.isEmpty m' then xs else Update m' : xs
  m' = M.mapMaybe simplOp m
simpl xs = xs

traceMove :: Int -> Int -> BigPass
traceMove o p r Nil = r :> Move (p - o)
traceMove o p r (Move n : xs) = traceMove o (p + n) r xs
traceMove o p r (Update m : xs) = traceMove o p (r :> Update (offsetUpd p m)) xs
traceMove o p r (Loop x : xs) = traceMove o p (r :> Loop x') xs
  where
  x' = { idx: x.idx + p, proc: run (traceMove p p) x.proc }
traceMove o p r (In i : xs) = traceMove o p (r :> In (i + p)) xs
traceMove o p r (Out i : xs) = traceMove o p (r :> Out (i + p)) xs

loop :: Pass
loop (Loop { idx, proc: Update m : Nil } : xs)
  | Just (Add (-1)) <- M.lookup idx m
  , Just m' <- traverseWithIndex (prod idx) m = Update m' : xs
loop xs = xs

offset :: Int -> Pass
offset o (Loop x : xs) = Loop x { idx = x.idx + o } : xs
offset o (In i : xs) = In (i + o) : xs
offset o (Out i : xs) = Out (i + o) : xs
offset o (Update m : xs) = Update (offsetUpd o m) : xs
offset _ xs = xs

foldMapIdx :: forall m. Monoid m => (Int -> m) -> Proc -> m
foldMapIdx f = foldMap (foldMapIdxI f)

foldMapIdxI :: forall m. Monoid m => (Int -> m) -> Instr -> m
foldMapIdxI _ (Move _) = mempty
foldMapIdxI f (Update m) = foldMapWithIndex (\i x -> f i <> foldMapIdxOp f x) m
foldMapIdxI f (In i) = f i
foldMapIdxI f (Out i) = f i
foldMapIdxI f (Loop { idx, proc }) = f idx <> foldMapIdx f proc

foldMapIdxOp :: forall m. Monoid m => (Int -> m) -> Oper -> m
foldMapIdxOp _ (Add _) = mempty
foldMapIdxOp f (Set e) = foldMapIdxE f e

foldMapIdxE :: forall m. Monoid m => (Int -> m) -> Expr -> m
foldMapIdxE _ (ECst _) = mempty
foldMapIdxE f (EDyn i) = f i
foldMapIdxE f (ENeg e) = foldMapIdxE f e
foldMapIdxE f (EAdd a b) = foldMapIdxE f a <> foldMapIdxE f b
foldMapIdxE f (ESub a b) = foldMapIdxE f a <> foldMapIdxE f b
foldMapIdxE f (EMul a b) = foldMapIdxE f a <> foldMapIdxE f b

mapKV :: forall k1 k2 v1 v2. Ord k2 => (k1 /\ v1 -> k2 /\ v2) -> M.Map k1 v1 -> M.Map k2 v2
mapKV f m = M.fromFoldable $ f <$> (M.toUnfoldable m :: Array _)

prod :: Int -> Int -> Oper -> Maybe Oper
prod o i (Add x)
  | o == i = Just $ Set $ ECst 0
  | otherwise = Just $ Set $ EAdd (EDyn i) (EMul (ECst x) (EDyn o))
prod _ _ _ = Nothing

operExpr :: Int -> Oper -> Expr
operExpr _ (Set e) = e
operExpr i (Add v) = EAdd (ECst v) (EDyn i)

instE :: (Int -> Expr) -> Expr -> Expr
instE f = recExpr f'
  where
  f' (EDyn i) = f i
  f' e = e

offsetUpd :: Int -> M.Map Int Oper -> M.Map Int Oper
offsetUpd p = mapKV (bimap (add p) (offsetOp p))

offsetE :: Int -> Expr -> Expr
offsetE o = instE (EDyn <<< add o)

offsetOp :: Int -> Oper -> Oper
offsetOp 0 x = x
offsetOp _ x@(Add _) = x
offsetOp o (Set e) = Set $ offsetE o e

simplOp :: Oper -> Maybe Oper
simplOp (Add 0) = Nothing
simplOp x@(Add _) = Just x
simplOp (Set e) = Just $ Set $ simplE e

simplE :: Expr -> Expr
simplE = recExpr simplE'

simplE' :: Expr -> Expr
simplE' (EAdd (ECst 0) a) = a
simplE' (EAdd a (ECst 0)) = a
simplE' (EAdd (ENeg a) b) = ESub b a
simplE' (EAdd a (ENeg b)) = ESub a b
simplE' (EMul (ECst 1) a) = a
simplE' (EMul a (ECst 1)) = a
simplE' (EMul (ECst (-1)) a) = ENeg a
simplE' (EMul a (ECst (-1))) = ENeg a
simplE' e = e

recExpr :: (Expr -> Expr) -> Expr -> Expr
recExpr f (ENeg e) = f $ ENeg (recExpr f e)
recExpr f (EAdd a b) = f $ EAdd (recExpr f a) (recExpr f b)
recExpr f (ESub a b) = f $ ESub (recExpr f a) (recExpr f b)
recExpr f (EMul a b) = f $ EMul (recExpr f a) (recExpr f b)
recExpr f e = f e

optimize :: Pass
optimize =
  run (traceMove 0 0)
    >>> run (liftPass $ merge >>> simpl)
    >>> run (liftPass $ loop)
    >>> run (liftPass $ merge >>> simpl)

-- 优化后，生成wasm前需要做这个
normalize :: Pass
normalize proc = Move maxOfs : run (liftPass $ offset (-minOfs)) proc
  where
  (Min minOfs) = Min 0 <> foldMapIdx Min proc -- 避免出现负偏移，本就有的偏移不动
  (Max maxOfs) = Max 0 <> foldMapIdx Max proc -- 避免地址为负
